###
### !pip install ffmpeg gradio
###
###
import gradio as gr
import subprocess
import tempfile
import os
import torch
import torchaudio
from transformers import Wav2Vec2Processor, Wav2Vec2ForCTC
from AudioTranscriber import AudioTranscriber
import soundfile
import json
import traceback
import numpy as np

# 配置FFmpeg路径（如果需要）
# os.environ["PATH"] += os.pathsep + '/path/to/ffmpeg'

# 颜色定义
SUCCESS_COLOR = "#4CAF50"
ERROR_COLOR = "#F44336"
INFO_COLOR = "#2196F3"

def convert_m4a_to_wav(m4a_file):
    """将M4A文件转换为WAV格式"""
    if not m4a_file:
        return None, "请上传M4A文件"
    
    print(f"typeof m4a_file {type(m4a_file)}")
    
    # 创建临时文件
    with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as temp_wav:
        temp_wav_path = temp_wav.name
    
    try:
        # 使用FFmpeg进行转换
        subprocess.run(
            [
                "ffmpeg",
                "-i", m4a_file,  # 输入文件
                "-acodec", "pcm_s16le",  # 音频编码
                "-ar", "16000",  # 采样率
                "-y",  # 覆盖已存在文件
                temp_wav_path  # 输出文件
            ],
            check=True,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE
        )
        return temp_wav_path, f'<span style="color:{SUCCESS_COLOR}">转换成功</span>'
    except subprocess.CalledProcessError as e:
        return None, f'<span style="color:{ERROR_COLOR}">转换失败: {e.stderr.decode("utf-8")}</span>'
    except Exception as e:
        return None, f'<span style="color:{ERROR_COLOR}">发生错误: {str(e)}</span>'

# 音频预测函数
def predict_audio(wav_file):
    """使用模型对WAV文件进行预测"""
    if not wav_file:
        return "无法进行预测：缺少WAV文件或模型"
    
    try:
        # 创建转录器实例并加载模型
        transcriber = AudioTranscriber()
        if transcriber.load_model():
            # 准备输入数据
            speech_array, sampling_rate = soundfile.read(wav_file)
            print(f"len of speech is {len(speech_array)}")
            print(f"type of speech is {type(speech_array)}")
            # 检查是否为立体声数据
            if isinstance(speech_array[0], list):
                print("it is a stero audio")
            # 检查是否为立体声
            if len(speech_array.shape) > 1 and speech_array.shape[1] == 2:
                print('it is a stero audio file')
                # 方法一：转换为单声道
                speech_array = np.mean(speech_array, axis=1)
            # 假设最小长度为 1000000
            #min_length = 1000000
            #if len(speech_array) < min_length:
            #    padding = [0] * (min_length - len(speech_array))
            #    speech_array = speech_array + padding
            json_request_data = {
                "speech_array": speech_array.tolist(),
                "sampling_rate": sampling_rate
            }
            json_request_str = json.dumps(json_request_data)
            
            # 执行预测
            result = transcriber.predict(json_request_str)
            if result:
                print(f"转录结果: {result[0]}")
            else:
                print("预测失败")
        else:
            print("模型加载失败，无法进行预测")
        
        return f'<span style="color:{INFO_COLOR}">预测结果:</span> {result[0]}'
    except Exception as e:
        traceback.print_exc()
        return f'<span style="color:{ERROR_COLOR}">预测失败: {str(e)}</span>'

# 创建Gradio界面
def create_app():
    """创建Gradio应用"""
    
    with gr.Blocks(title="音频处理与预测系统", theme=gr.themes.Soft()) as app:
        gr.Markdown("## M4A转WAV并使用AI模型预测")
        gr.Markdown("上传M4A文件，系统将自动转换为WAV格式并使用AI模型进行预测。")
        
        with gr.Row():
            with gr.Column(scale=1):
                input_audio = gr.Audio(label="上传M4A文件", type="filepath" )
                convert_btn = gr.Button("转换为WAV", variant="primary")
                
                with gr.Row():
                    status_text = gr.HTML(label="状态")
                
                converted_audio = gr.Audio(label="转换后的WAV文件", type="filepath")
                predict_btn = gr.Button("运行预测", variant="secondary")
            
            with gr.Column(scale=1):
                output_text = gr.HTML(label="预测结果")
        
        # 设置事件处理
        convert_btn.click(
            fn=convert_m4a_to_wav,
            inputs=[input_audio],
            outputs=[converted_audio, status_text]
        )
        
        predict_btn.click(
            fn=lambda wav: predict_audio(wav),
            inputs=[converted_audio],
            outputs=[output_text]
        )
        
        # 应用启动时的提示
        gr.Markdown("""
        ### 使用说明
        1. 上传M4A格式的音频文件
        2. 点击"转换为WAV"按钮进行格式转换
        3. 转换完成后，点击"运行预测"按钮使用AI模型进行分析
        4. 预测结果将显示在右侧面板中
        
        *注意：当前示例使用的是语音识别模型，会将音频内容转换为文本。你可以替换为自己的分类模型以实现特定的音频分类任务。*
        """)
    
    return app

# 启动应用
if __name__ == "__main__":
    app = create_app()
    app.launch()
